package com.xjtu.util;

import com.xjtu.model.Topology;
import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.command.CreateContainerResponse;
import com.github.dockerjava.api.model.ContainerNetwork;
import com.github.dockerjava.api.model.Network;
import com.github.dockerjava.core.DockerClientBuilder;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

@Slf4j
@Component
public class DockerUtil {
    private static final DockerClient dockerClient = DockerClientBuilder.getInstance("tcp://localhost:2375").build();

    /**
     * 根据节点类型选择镜像创建容器
     */
    public static String createContainer(String containerName, String type) {
        String imageName = getImageByType(type);
        CreateContainerResponse container = dockerClient.createContainerCmd(imageName)
                .withName(containerName)
                .exec();
        dockerClient.startContainerCmd(container.getId()).exec();
        log.info("容器创建成功，容器ID: {}", container.getId());
        return container.getId();
    }

    /**
     * 批量创建容器
     */
    public static Map<String, String> createContainers(List<Topology.Node> nodes) {
        Map<String, String> containerIdMap = new HashMap<>();
        for (Topology.Node node : nodes) {
            String containerId = createContainer(node.getName(), node.getType());
            containerIdMap.put(node.getName(), containerId);
            log.info("容器name: {},容器ID: {}",node.getName(), containerId);
        }
        return containerIdMap;
    }

    /**
     * 获取节点类型对应的镜像
     */
    private static String getImageByType(String type) {
        switch (type) {
            case "attacker":
                return "attacker-image:latest"; // 攻击节点镜像
            case "server":
                return "server-image:latest";  // 服务器节点镜像
            default:
                throw new IllegalArgumentException("Unknown node type: " + type);
        }
    }

    /**
     * 将容器连接到指定网络，并分配 IP 地址
     */
    public static void connectToNetwork(String containerId, String networkName, String ipAddress) {
        Optional<Network> existingNetwork = dockerClient.listNetworksCmd()
                .exec()
                .stream()
                .filter(net -> net.getName().equals(networkName))
                .findFirst();

        if (existingNetwork == null) {
            dockerClient.createNetworkCmd()
                    .withName(networkName)
                    .withIpam(new Network.Ipam().withConfig(new Network.Ipam.Config().withSubnet("192.168.1.0/24")))
                    .exec();
        }

        dockerClient.connectToNetworkCmd()
                .withContainerId(containerId)
                .withContainerNetwork(new ContainerNetwork().withIpv4Address(ipAddress))
                .withNetworkId(existingNetwork.orElseGet(() -> dockerClient.listNetworksCmd()
                        .exec()
                        .stream()
                        .filter(net -> net.getName().equals(networkName))
                        .findFirst()
                        .orElseThrow(() -> new RuntimeException("Failed to create or find network"))
                ).getId())
                .exec();
    }

    /**
     * 批量将容器连接到网络
     */
    public static void connectContainersToNetwork(Map<String, String> containerIdMap, String networkName, Map<String, String> ipAddressMap) {
        for (Map.Entry<String, String> entry : containerIdMap.entrySet()) {
            String containerName = entry.getKey();
            String containerId = entry.getValue();
            String ipAddress = ipAddressMap.get(containerName);
            connectToNetwork(containerId, networkName, ipAddress);
        }
    }

    /**
     * 删除单个容器
     */
    public static void removeContainer(String containerId) {
        dockerClient.removeContainerCmd(containerId).withForce(true).exec();
    }

    /**
     * 批量删除容器
     */
    public static void removeContainers(List<String> containerIds) {
        for (String containerId : containerIds) {
            removeContainer(containerId);
        }
    }
}
